package com.kklsqm.langchain4j.service;

import cn.hutool.json.JSONObject;
import cn.hutool.json.JSONUtil;
import okhttp3.*;
import org.springframework.stereotype.Service;

import javax.imageio.ImageIO;
import java.awt.image.BufferedImage;
import java.io.ByteArrayInputStream;
import java.io.IOException;
import java.util.Base64;

@Service
public class ImageEraseCompletionService {

    private static final String API_URL = "https://dashscope.aliyuncs.com/api/v1/services/aigc/image2image/image-synthesis";
    private static final String QUERY_URL_TEMPLATE = "https://dashscope.aliyuncs.com/api/v1/tasks/%s";
    private static final String API_KEY = "sk-072acf6e9113491bb5809110bd83e8b1"; // 替换成你的真实Key

    private final OkHttpClient client = new OkHttpClient();

    public String eraseImage(String imageUrl, String maskUrl,String foregroundUrl) {
        try {
            // 构造请求体
            JSONObject input = new JSONObject()
                    .set("image_url", imageUrl)
                    .set("mask_url", maskUrl)
                    .set("foreground_url", foregroundUrl);

            JSONObject parameters = new JSONObject()
                    .set("dilate_flag", true)
                    .set("fast_mode", true)
                    .set("add_watermark", false);

            JSONObject body = new JSONObject()
                    .set("model", "image-erase-completion")
                    .set("input", input)
                    .set("parameters", parameters);

            // Step 1: 创建任务
            Request request = new Request.Builder()
                    .url(API_URL)
                    .addHeader("Authorization", "Bearer " + API_KEY)
                    .addHeader("Content-Type", "application/json")
                    .addHeader("X-DashScope-Async", "enable")
                    .post(RequestBody.create(body.toString(), MediaType.parse("application/json")))
                    .build();

            try (Response response = client.newCall(request).execute()) {
                if (!response.isSuccessful()) {
                    String errorBody = response.body().string();
                    throw new IOException("创建任务失败，HTTP状态码: " + response.code() + "\n错误信息: " + errorBody);
                }

                JSONObject json = JSONUtil.parseObj(response.body().string());
                String taskId = json.getJSONObject("output").getStr("task_id");

                // Step 2: 轮询任务结果
                return pollTaskResult(taskId);
            }
        } catch (Exception e) {
            throw new RuntimeException("调用图像擦除补全服务失败: " + e.getMessage(), e);
        }
    }

    private String pollTaskResult(String taskId) throws IOException, InterruptedException {
        String queryUrl = String.format(QUERY_URL_TEMPLATE, taskId);
        Request queryRequest = new Request.Builder()
                .url(queryUrl)
                .addHeader("Authorization", API_KEY)
                .get()
                .build();

        while (true) {
            try (Response queryResponse = client.newCall(queryRequest).execute()) {
                if (!queryResponse.isSuccessful()) {
                    throw new IOException("查询任务状态失败，HTTP状态码: " + queryResponse.code());
                }

                JSONObject result = JSONUtil.parseObj(queryResponse.body().string());
                JSONObject output = result.getJSONObject("output");
                String status = output.getStr("task_status");

                System.out.println("当前任务状态: " + status);

                if ("SUCCEEDED".equals(status)) {
                    return output.getStr("output_image_url");
                } else if ("FAILED".equals(status)) {
                    throw new RuntimeException("图像擦除任务执行失败: " + result.toStringPretty());
                }

                Thread.sleep(3000); // 等待3秒后重试
            }
        }
    }

    // 将 Base64 字符串解码为 BufferedImage
    public BufferedImage decodeBase64ToImage(String base64Str) throws IOException {
        // 去除 data:image/png;base64, 前缀（如果有）
        String base64Data = base64Str.contains(",") ? base64Str.split(",")[1] : base64Str;
        byte[] imageBytes = Base64.getDecoder().decode(base64Data);
        return ImageIO.read(new ByteArrayInputStream(imageBytes));
    }

    //  验证图像尺寸
    public void validateImageSize(BufferedImage image) {
        int width = image.getWidth();
        int height = image.getHeight();
        if (width < 512 || height < 512 || width > 4096 || height > 4096) {
            throw new IllegalArgumentException("图像尺寸非法，宽高必须在 512 到 4096 像素之间！");
        }
    }

    // 验证图像文件大小
    public void validateImageFileSize(String base64Str) {
        String base64Data = base64Str.contains(",") ? base64Str.split(",")[1] : base64Str;
        int byteSize = Base64.getDecoder().decode(base64Data).length;
        if (byteSize > 10 * 1024 * 1024) {
            throw new IllegalArgumentException("图像文件太大，超过 10MB");
        }
    }

    //  验证图像格式
    public void validateBase64Image(String base64Str) throws IOException {
        validateImageFileSize(base64Str);
        BufferedImage image = decodeBase64ToImage(base64Str);
        if (image == null) {
            throw new IllegalArgumentException("无法解析图像，格式可能不正确");
        }
        validateImageSize(image);
    }
}
